import argparse
from copy import deepcopy
from pathlib import Path

import yaml

import utils


def run():
    # Valid argument
    opts = vars(opt)
    assert sum(opts[k] for k in opts) == 1, 'ERROR - you must specify only and at least 1 argument'

    # Get the action to proceed
    action = ''
    for k, v in opts.items():
        if v:
            action = k
            break

    # Config yaml
    root = Path('.')
    config = yaml.load(open(root / 'config' / (action + '.yaml')), Loader=yaml.FullLoader)
    print(config)
    config = deepcopy(config)

    if opt.select:
        dir_ = config['dir']
        src_labels, src_images, target_labels, target_images = iter(map(lambda x: Path(dir_[x]), dir_))
        utils.select(src_labels, src_images, target_labels, target_images, config['num'], config['classes'])
    elif opt.split:
        dir_ = config['dir']
        src, target = iter(map(lambda x: Path(dir_[x]), dir_))
        utils.split(src, target, config['train'], config['valid'], config['test'])
    elif opt.intersect:
        dir_ = config['dir']
        src_labels, src_images, target_labels, target_images = iter(map(lambda x: Path(dir_[x]), dir_))
        utils.intersect(src_labels, src_images, target_labels, target_images)
    elif opt.rename:
        utils.rename(Path(config['src']), config['basename'])
    else:
        raise NotImplementedError('Not in opt dict')


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--select', action='store_true', help='select specific classes from given ones')
    parser.add_argument('--split', action='store_true', help='split dataset to train, valid, test')
    parser.add_argument('--intersect', action='store_true', help='get intersection of images and labels on basename')
    parser.add_argument('--rename', action='store_true', help='rename all files in a folder')

    opt = parser.parse_args()
    print(opt)

    run()
